package standford;

import java.util.*;

import ai.graph.NamedRelation;
import ai.graph.WordNode;

import edu.stanford.nlp.trees.*;
import edu.stanford.nlp.ling.MapLabel;
import edu.stanford.nlp.parser.lexparser.LexicalizedParser;
import edu.uci.ics.jung.graph.Graph;
import edu.uci.ics.jung.graph.Vertex;
import edu.uci.ics.jung.graph.impl.DirectedSparseEdge;
import edu.uci.ics.jung.graph.impl.DirectedSparseGraph;
import edu.uci.ics.jung.graph.impl.DirectedSparseVertex;
import graph.SimpleGraphDraw;

public class ParserDemo {
	public static void main(String[] args) {
		LexicalizedParser lp = new LexicalizedParser(
				"M:/Projects/java/_libs/stanford/stanford-parser-2007-08-19/englishPCFG.ser.gz");
		lp.setOptionFlags(new String[] { "-maxLength", "80",
				"-retainTmpSubcategories" });

		// String[] sent = { "This", "is", "an", "easy", "sentence", "." };
		String[] sent = { "A", "musician", "is", "a", "person", "who", "plays",
				"or", "composes", "music", "." };
		Tree parse = (Tree) lp.apply(Arrays.asList(sent));
		parse.pennPrint();
		System.out.println();

		TreebankLanguagePack tlp = new PennTreebankLanguagePack();
		GrammaticalStructureFactory gsf = tlp.grammaticalStructureFactory();
		GrammaticalStructure gs = gsf.newGrammaticalStructure(parse);
		Collection<TypedDependency> tdl = gs.typedDependenciesCollapsed();
		System.out.println(tdl);
		System.out.println();

		TreePrint tp = new TreePrint("penn,typedDependenciesCollapsed");
		tp.printTree(parse);
		SimpleGraphDraw.showGraph(processDependency(gs));
	}

	/*
	 * det(musician-2, A-1) nsubj(person-5, musician-2) cop(person-5, is-3)
	 * det(person-5, a-4) nsubj(plays-7, who-6) rcmod(person-5, plays-7)
	 * cc(plays-7, or-8) conj(plays-7, composes-9) dobj(plays-7, music-10)
	 */
	public static final Set<String> ACCEPTED = new HashSet<String>(Arrays
			.asList("nsubj", "conj", "dobj", "cc", "cop", "rcmod"));

	public static final Set<String> SKIP = new HashSet<String>(Arrays
			.asList("det"));

	public static Graph processDependency(GrammaticalStructure s) {
		Graph g = new DirectedSparseGraph();
		Collection<TypedDependency> tdl = s.typedDependenciesCollapsed();
		Map<TreeGraphNode, WordNode> set = new HashMap<TreeGraphNode, WordNode>();
		for (TypedDependency typedDep : tdl) {
			// skip some relation types
			if (SKIP.contains(typedDep.reln().getShortName())) {
				continue;
			}

			WordNode dep = set.get(typedDep.dep());
			WordNode gov = set.get(typedDep.gov());
			if (dep == null) {
				MapLabel label = (MapLabel) typedDep.dep().label();
				dep = (WordNode) g.addVertex(new WordNode(label.word(), label
						.tag()));
				set.put(typedDep.dep(), dep);

			}
			if (gov == null) {
				MapLabel label = (MapLabel) typedDep.gov().label();
				g.addVertex(new WordNode(label.word(), label.tag()));
				gov = (WordNode) g.addVertex(new WordNode(label.word(), label
						.tag()));
				set.put(typedDep.gov(), gov);
			}

			// create relations
			String relName = typedDep.reln().getShortName();
			if (relName.startsWith("nsubj")) {
				g.addEdge(new NamedRelation(dep, gov, relName));
			} else {
				g.addEdge(new NamedRelation(gov, dep, relName));
			} // TODO: conj
		}

		return g;
	}

}